Megatron-Core Custom FSDP

Jianbin Chang, Hongbin Liu

议程
全分片数据并行（FSDP）简介
- FSDP 工作原理：ZeRO-DP 分片策略
- FSDP 工作原理：FSDP 单元 (FSDP-Unit
- FSDP 工作原理：AllGather
- FSDP 工作原理：Reduce-Scatter
- FSDP 工作原理：不同 ZeRO-DP 策略中的 AllGather (AG) 与 Reduce-Scatter (RS
自定义 FSDP - 关键特性
- Megatron-Core 自定义 FSDP 关键特性架构图
- 并行通信对 Gemm 的不利影响
- 基于 UBR 的 SM 高效 NCCL 通信
- 基于 UBR 的 SM 高效 NCCL 通信：AllGather 示例
- NCCL 组调用 (Group Calls
理论分析：FSDP 中的通信隐藏
- FSDP 中的计算与通信成本：Transformer 层前向传播
- FSDP 中的计算与通信成本：Transformer 层后向传播
- FSDP 中的计算与通信成本：MoE 层后向传播
- FSDP 中的计算与通信成本：调优并行策略
基准测试与案例研究
- 来自 BioNeMo 的自定义 FSDP 反馈
- 基准测试：vs. FSDP2
- 基准测试：vs. 3D 并行
- FP8 训练中的挑战：CPU 瓶颈问题及 FSDP 解决方案
- 案例研究：Llama3-70B + FSDP
最佳实践
如何在 Megatron-Core 中启用自定义 FSDP

议程

FSDP 简介
自定义 FSDP - 关键特性
理论分析：FSDP 中的通信隐藏
基准测试与案例研究
最佳实践

全分片数据并行（FSDP）简介

什么是 FSDP？
- 定义：FSDP 是一种数据并行训练技术，它将模型参数、梯度和优化器状态分片到多个 GPU 上，从而减少内存使用，使得更大模型的训练成为可能。
- 比较：与在多个 GPU 间复制模型状态的传统数据并行方法不同，FSDP 对这些状态进行分片，并且还可以将它们卸载到 CPU。
FSDP 的优势
- 可扩展性：能够训练因内存限制而无法实现的非常大的模型。
- 灵活性：支持各种模型架构和优化器，适用于广泛的 AI 应用。
- 性能：通过高效管理 GPU 资源，减少内存占用并提高训练速度。
FSDP 的工作原理
- ZeRO-DP 分片策略
- FSDP 单元 (FSDP Unit)
- 参数 AllGather 和梯度 Reduce-Scatter

参考 PyTorch FSDP 介绍：https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html

FSDP 工作原理：ZeRO-DP 分片策略

FSDP 的分片策略能够显著减少每个 GPU 的内存消耗，但更激进的分片会增加通信开销。

基线 (Baseline)：不使用分片，内存消耗为 (2 + 2 + K) * Ψ，在此示例中为 120GB。通信开销为 1 次 all-reduce。
P_os (优化器状态分片)：对优化器状态进行分片，内存消耗降至 31.4GB。通信开销为 1 次 all-reduce。
P_os+g (优化器状态+梯度分片)：进一步对梯度进行分片，内存消耗降至 16.6GB。通信开销变为 1 次 all-gather 和 GA 次 reduce-scatter。
P_os+g+p (优化器状态+梯度+参数分片)：对所有状态进行分片，内存消耗降至 1.9GB。通信开销为 2GA 次 all-gather 和 GA 次 reduce-scatter。

FSDP 工作原理：FSDP 单元 (FSDP-Unit)

FSDP 单元是我们用于取消分片/重新分片参数的操作单元，确保这些参数具有相同的计算生命周期。

模型的层（如 layer0 到 layer5）被包装并分片到不同的 FSDP 单元中（如 FSDP Unit0, Unit1, Unit2）。
前向传播：在执行计算前，通过 gather full params 操作收集完整的参数。计算后，通过 free peer shards 释放从其他 GPU 获取的参数分片。
反向传播：在反向传播期间，首先同步梯度，然后释放参数分片。接着为下一层收集参数。

FSDP 工作原理：AllGather

分片 (Sharding)：一个 FSDP 单元内的权重（W1, W2, W3, W4）被分片到四个不同的节点上。
All-Gather：在进行前向或反向传播之前，每个节点会从其他所有节点收集它们所持有的权重分片，从而在本地重构出完整的权重。例如，Node1 在收集后会拥有 W1, W2, W3, W4。
释放对等分片 (Free Peer Shards)：在使用完毕后，每个节点会释放掉从其他节点收集来的权重分片，只保留自己负责的那一部分（例如 Node1 只保留 W1）。

FSDP 工作原理：Reduce-Scatter

All-Gather：与前述相同，在计算前，所有节点都拥有完整的权重。
反向传播 (Loss.backward())：计算损失并进行反向传播后，每个节点都会计算出针对所有权重的梯度（例如 Node1 计算出 G11, G12, G13, G14）。
Reduce-Scatter：该操作会将所有节点上的梯度进行聚合（相加），然后将聚合后的结果分发回各个节点。每个节点只接收自己负责的那一部分梯度（例如 Node1 接收 ΣG_i1）。

FSDP 工作原理：不同 ZeRO-DP 策略中的 AllGather (AG) 与 Reduce-Scatter (RS)

下图展示了在不同分片策略下，通信操作（AG 和 RS）与计算操作（前向 F、反向 B、优化器步骤 Opt step）的穿插执行方式。

P os (优化器状态分片): RS 操作在反向传播和优化器步骤之间进行。AG 操作在前向传播之前进行。
P os+g (优化器状态+梯度分片): RS 操作在反向传播期间穿插进行。
P os+g+p (优化器状态+梯度+参数分片): AG 和 RS 操作在整个前向和反向传播过程中都穿插进行。

自定义 FSDP - 关键特性

我们已经支持在 M-Core 上使用 Torch FSDP2。M-Core Custom FSDP 是一个轻量级、面向性能的 FSDP 实现。

下表比较了 Torch FSDP2 和我们的 Custom FSDP：

特性	FSDP2 (Torch)	Custom FSDP
ZeRO-DP 等效性	✅	✅
可扩展性	✅	✅✅
性能	✅	✅✅
内存管理	✅	✅✅ (可选的双缓冲分配器)
CPU 卸载	✅	❌
通信定制	❌	✅ (可选的基于用户缓冲区注册的 NCCL 通信)
FP8 训练	✅	✅✅ (原生 FP8 支持；更好的转置缓存管理；NCCL 组调用)
M-Core MoE 训练兼容	❌	✅

FSDP2(Torch) 与 Custom FSDP 功能对比表，Page 10

Megatron-Core 自定义 FSDP 关键特性架构图

下图展示了 Custom FSDP 的核心组件及其交互关系。

FullyShardedDataParallel：顶层模块，负责动态删除 fp8 转置缓存，并在前向/后向 hook 中添加必要的原生 fp8 量化过程。
AllGatherPipeline / GradReducePipeline：实现计算与通信的重叠，并使用 NCCL 组调用来提高效率。
ParamAndGradBuffer：管理参数和梯度的缓冲区。
DataParallelBuffer / TemporaryBucketAllocator：支持基于用户缓冲区注册的 NCCL 通信，并使用双缓冲分配器。
torch.nn.Module：通过 Per-bucket NCCL communication: ZeRO-Copy 进行参数桶映射。

Megatron-Core Custom FSDP 关键特性架构图，Page 11

并行通信对 Gemm 的不利影响

通信流处理器（SM）对 Gemm 性能的影响

当 Gemm（通用矩阵乘法）和 Comm（通信）操作并发运行时，它们必须共享资源（SM、DRAM、L2 缓存等），这会导致性能下降。

Gemm 在与 Comm 重叠时吞吐量较低的原因：
1. Comm 从 Gemm 中占用 SM（SM 干扰）。
2. Comm 从 Gemm 中占用 DRAM/L2 带宽（内存干扰）。
LLM 训练中的 Gemm 对 SM 干扰比内存干扰更敏感。

下图展示了 Gemm 性能在减少可用内存带宽（左）和减少 SM 数量（右）时的下降情况。

基于 UBR 的 SM 高效 NCCL 通信

利用 NVL Sharp、IB-Sharp 和 NCCL 实现高效集合通信

NVLink Sharp 和 IB-Sharp 是硬件加速的集合通信技术，可将归约操作卸载到网络交换机。
NCCL v2.27 集成了这些算法，能够使用更少的 GPU SM 资源来执行高带宽的 all-gather 和 reduce-scatter 操作。
传统算法在 FSDP 训练中会消耗更多的 SM，从而降低计算效率。借助新的 NCCL 算法，FSDP 可以实现更好的性能。

下图展示了在 256 个 GPU（8x32）上，不同 NCCL 算法在 All-Gather (AG) 和 Reduce-Scatter (RS) 操作中的带宽表现。新的 NVLS+IB/S 算法（黑色曲线）使用极少的 CTA（4个）即可达到高带宽。
*预计在 Megatron-Core v0.13 中发布。

基于 UBR 的 SM 高效 NCCL 通信：AllGather 示例

什么是 UBR？

Torch 原生 AllGather:
1. 使用 Torch 内存分配器分配输入和输出张量。
2. 在 Torch 原生张量上执行 all-gather。
基于 UBR 的 AllGather:
1. 使用 NCCL 内存分配器分配输入和输出张量。
2. 在已注册的用户缓冲区张量上执行 all-gather。

目前，基于用户缓冲区注册（UBR）的 NCCL 通信只能由带有双缓冲内存分配器的 Custom FSDP 支持。

Torch 原生 AllGather 与 UBR-based AllGather 对比图，Page 14

NCCL 组调用 (Group Calls)

用于加速多桶通信的 NCCL 组调用

在实际的 FSDP 实现中，由于混合通信精度和参数映射限制等约束，常常需要将单个 FSDP 单元拆分为多个通信组。
每个组都会触发一次单独的通信核函数启动，导致效率低下。

示例：单个 FSDP 单元被拆分为三个通信组。

下图高亮部分显示，由于数据类型（bfloat16, float8）和模块的不同，参数被分到了多个通信组（group 14, 15, 16, 17）中，这会降低通信效率。

问题：当一个 FSDP 单元由多个受延迟限制的通信组组成时，通信效率会很差。
解决方案：通过 NCCL 组调用，这些受延迟限制的通信核心（communication kernels）可以被融合（fused），从而提高通信效率。

上图对比了未使用（上）和使用（下）聚合 NCCL 通信的训练过程。在不使用聚合通信时，每个 FSDP 单元会为 FP8 和 BF16 的桶（bucket）分别触发大量独立的 NCCL 核心。使用聚合 NCCL 通信后，每个 FSDP 单元仅启动一个 NCCL 核心，显著减少了核心启动的开销，从而将通信操作融合，提高了效率。

理论分析：FSDP 中的通信隐藏

FSDP 中的计算与通信成本：Transformer 层前向传播

设置: 1个 Transformer 层，包含多头注意力（MHA）和 GeLU 激活函数。
前向传播计算量 (V_comp):
- 公式：V_comp = 2 x (3h^2+4h^2+4h^2)/tp x SL/cp x mbs，其中 2 表示每次参数操作对应 2 个浮点运算（fma ops）。
通信量 (V_comm):
- 公式：V_comm = 2 x (3h^2+4h^2+4h^2)/tp，其中 2 表示参数存储于 bf16 格式。
完全隐藏通信的条件:
- t_comp > t_comm，即计算时间大于通信时间。
- 推导可得：tokens_per_micro_batch_per_gpu > (achieved Flops / achieved BW)
变量定义:
- h: 隐藏层维度 (hidden dimension)
- SL: 序列长度 (sequence length)
- mbs: 微批次大小 (micro batch size)
- tp: 张量并行大小 (tensor parallel size)
- cp: 上下文并行大小 (context parallel size)

此页在上一页的分析基础上，增加了一个表格，展示了在不同张量并行（TP）规模下，为隐藏通信所需的最小微批次 token 数。

结论: 张量并行（TP）通信会与 FSDP 的通信竞争 NVLink 带宽。因此，我们需要调整 TP 和 FSDP 的规模以最大化性能。
数据表分析:
- 系统配置：H100, 8*400Gbs Infiniband，效率为 85%。
- 随着 TP 规模从 1 增加到 8，实现的带宽（Achieved BW）从 340 GB/s 降低到 42.5 GB/s。
- 为了完全隐藏通信，所需的最小微批次 token 数也随之增加，从 1.8K 增加到 14.1K。

FSDP 中的计算与通信成本：Transformer 层后向传播

设置: 与前向传播相同。
后向传播计算量 (V_comp):
- 公式：V_comp = 4 x (3h^2+4h^2+4h^2)/tp x SL/cp x mbs，其中 4 表示每次参数操作对应 4 个浮点运算（fma ops）。
通信量 (V_comm):
- 公式：V_comm = (2+4) x (3h^2+4h^2+4h^2)/tp，其中 2 和 4 分别表示参数和梯度存储于 bf16 和 fp32 精度。
完全隐藏通信的条件:
- 推导可得：tokens_per_micro_batch_per_gpu > (3/2 * achieved Flops / achieved BW)
数据表分析:
- 在后向传播中，隐藏通信所需的最小微批次 token 数比前向传播更高，从 TP=1 时的 2.6K 增加到 TP=8 时的 21.2K。
结论:
- 在后向传播中，FSDP 的通信更难被完全隐藏。
- 将梯度存储为 bf16 格式有助于减轻通信成本。

FSDP 中的计算与通信成本：MoE 层后向传播

设置: 1个包含 N 个专家的 MoE 层。
后向传播计算量 (V_comp):
- 公式：V_comp = 4 x (4h^2+4h^2)/tp x SL/cp x mbs x topK，其中 4 表示 fma ops。
通信量 (V_comm):
- 公式：V_comm = (2+4) x (4h^2+4h^2)/tp x N/ep，其中 2 和 4 分别表示参数和梯度存储于 bf16 和 fp32 精度。
完全隐藏通信的条件:
- 推导可得：tokens_per_micro_batch_per_gpu > (3/2 * N/ep * 1/topK * achieved Flops / achieved BW)
结论:
- 大量专家会增加 FSDP 通信暴露的可能性。
- 结合大的专家并行规模（expert parallel size）和 topK 值，可以大幅降低 FSDP 中的通信与计算比率，从而使 FSDP 中的通信能被很好地重叠（overlapped）。
新增变量:
- topK: 每个 token 被分派到的专家数量
- N: 专家总数

FSDP 中的计算与通信成本：调优并行策略

本页总结了如何调整并行策略以更好地隐藏通信成本。

判据: tokens_per_micro_batch_per_gpu > (3/2 * N/ep * 1/topK * achieved Flops / achieved BW)
FSDP 通信暴露较少的条件:
- 长上下文长度 (SL)
- 大的微批次大小 (mbs)
- 小的上下文并行规模 (cp)
- 大的专家并行规模 (ep)
- 大的 topK 值
- 计算能力较低的 GPU (achieved Flops)
- 具有更快节点内/节点间带宽的 GPU (achieved BW)
- 小的张量并行规模 (achieved BW)

基准测试与案例研究

来自 BioNeMo 的自定义 FSDP 反馈

该基准测试在 2 个节点、16xH100 GPU 上进行。

测试结果摘要:
- ESM-2 3B: 自定义 FSDP (cFSDP) 与 Zero 1 (HF Accelerate) 取得了相似的模型浮点运算利用率（MFU%），均为 53.5%。
- ESM-2 15B:
  - 使用 FSDP2 (HF Accelerate, FA2)，MFU% 为 35.3%。
  - 使用 cFSDP (批次大小为4)，MFU% 达到 49.4%，性能提升 85.2%。
  - 使用 cFSDP (批次大小为6)，MFU% 达到 53.3%，性能提升 99.6%。

基准测试：vs. FSDP2

此图展示了在 Megatron-Core 上训练 LLaMA3-70B 时，自定义 FSDP (C-FSDP) 与 FSDP2 的扩展效率对比。

性能对比 (PFlops): C-FSDP 的聚合吞吐量（PetaFLOP/s）始终高于 FSDP2，并且随着 GPU 数量的增加，性能优势从 1.06 倍扩大到 1.12 倍。
内存效率对比: C-FSDP 相对于 FSDP2 的内存节约比例也随着规模扩大而增加，从 1.05 倍提升至 1.76 倍。
结论: 自定义 FSDP (Custom FSDP) 在性能、扩展效率和内存节省能力方面均优于 FSDP2。

基准测试：vs. 3D 并行

此图比较了自定义 FSDP (C-FSDP) 与 3D 并行在 H100 系统上的性能和内存使用情况。

性能 (TFLOPS/GPU): C-FSDP 的性能与 3D 并行相当。
- Llama3-70B BF16: C-FSDP (427 TFLOPS) vs. 3D-Parallel (465 TFLOPS)
- Llama3-70B FP8: C-FSDP (660 TFLOPS) vs. 3D-Parallel (652 TFLOPS)
内存使用 (GB): C-FSDP 在所有测试场景下都比 3D 并行占用更少的内存。
- Llama3-70B FP8: C-FSDP (65.63 GB) vs. 3D-Parallel (69.6 GB)
结论: 在经典工作负载上，自定义 FSDP 的性能与 3D 并行相当，并具有内存节省的优势。

FP8 训练中的挑战：CPU 瓶颈问题及 FSDP 解决方案

问题描述: CPU 瓶颈
当使用 FP8 精度时，GPU 执行速度变得非常快，以至于超过了 CPU 启动 GPU 核心（kernel）的速度，导致 GPU timeline 上出现空闲气泡（launch bubbles），形成 CPU 瓶颈。下图的 profiler timeline 显示了 Llama3-8B 模型在使用 3D 并行和 FP8 精度时，GPU 计算流之间存在明显的间隙，标注为“CPU Launch Bound”。

解决方案: FSDP
FSDP 通过分散参数、梯度和优化器状态来节省内存。这使得可以使用更大的 mbs x seq-length（微批次大小 x 序列长度），增加了单个核心的计算负载，从而避免了因核心启动过快而产生的气泡问题。
- 上图 (3D 并行): Llama3-8B FP8, CP2DP4 的 timeline 显示了大量细碎的计算核心。
- 下图 (FSDP): Llama3-8B FP8, FSDP 的 timeline 显示了更大、更连续的计算核心，有效地缓解了 CPU 启动瓶颈。

案例研究：Llama3-70B + FSDP

挑战：内存溢出（OOM）

问题场景: 使用 FSDP64（64 位优化器状态）且不带激活重计算（activation recompute）时，发生内存溢出（OOM）。
内存分析:
- 激活内存（Activation Memory）占用超过 70GB。
- 模型权重和梯度（Model weights & Main Gradients）占用约 10GB。
根本原因: 巨大的激活内存导致显存耗尽。堆栈跟踪信息显示，内存分配峰值发生在 transformer 的前向传播过程中。
目标与挑战: 移除激活重计算是为了获得性能增益，但这直接导致了内存溢出问题。

解决方案 1：结合上下文并行（CP）

方法: 采用 FSDP64 + CP8（上下文并行度为8）的配置成功运行。
结果:
- 实现了 225.59 TFLOPS/GPU 的性能。
- 峰值内存占用得到有效控制。
内存分析: 通过上下文并行，激活内存（Activation Memory）被分散到多个 GPU 上，单个 GPU 的激活内存占用降低到约 25GB，从而避免了 OOM。下图展示了使用 FSDP64 + CP4 配置时的内存使用快照，激活内存峰值仍超过50GB，说明调整CP值对内存控制至关重要。

Page 31 展示了 Llama3-70B + FSDP 训练过程中的内存使用时间线。峰值激活内存超过50GB，而优化器状态占用约25GB。

结论: 结合 FSDP 和上下文并行（CP），可以在不使用激活重计算的情况下成功训练大模型，平衡了性能和内存消耗。

性能对比：上下文并行 vs. 激活重计算

本节通过表格和性能分析器视图，对比了采用上下文并行（CP）与采用激活重计算（Activation Recompute）两种方案的性能差异。

实验设置:
- 模型: Llama3-70b
- 计算精度: BF16
- GPU数量: 64
- 序列长度 (SEQ_LEN): 8194
- FSDP 并行度: 64
- 张量并行度 (TP): 1
- 流水线并行度 (PP): 1
- 微批次大小 (MBS): 1
- 全局批次大小 (GBS): 128

对比结果:
下表对比了三种不同配置下训练 Llama3-70B 的性能指标。

Page 32 的表格对比了三种不同配置下训练 Llama3-70B 的性能指标。

配置	激活完全重计算	内存使用 (GB)	TFLOPS/GPU
CP=2 + 激活重计算	是	64.62	427
CP=8 (无重计算)	否	79.05	226
CP=4 (无重计算)	否	79.05	379

性能分析:
- 上图 (FSDP + CP4): 时间线显示计算（Kernel）和通信（蓝色条块）之间存在明显的空闲间隙。标注指出：“计算无法隐藏通信”，这意味着 GPU 在等待通信完成，导致计算单元闲置，降低了效率。
- 下图 (FSDP + 激活重计算): 计算流（Compute Stream）变得非常繁忙和密集。标注指出：“完全隐藏通信”和“繁忙的计算流”。这表明通过重计算引入的额外计算任务成功地与通信操作重叠，使得 GPU 能够持续执行计算任务，从而提高了整体的吞吐量。

结论:
在此场景下，激活重计算是最佳的映射方案。尽管激活重计算增加了约 1/3 的计算量，但通过更好地重叠计算和通信，获得了更高的整体性能，同时显著降低了内存使用。

最佳实践

本页总结了在训练大模型时的一些最佳实践建议。

实践领域	建议
计算与通信重叠	使用大的序列长度/微批次大小；调整通信桶（communication bucket）大小和 SM 分配。
激活内存优化	根据需要使用上下文并行（context parallelism）、激活卸载（activation offloading）或选择性激活重计算（selective activation recompute）。
权衡：内存节省 vs. 通信	仅在需要时启用重计算；调整重计算的粒度；优先考虑重叠以提高吞吐量。

如何在 Megatron-Core 中启用自定义 FSDP

本页提供了在 Megatron-Core 框架中启用和配置自定义 FSDP（Fully Sharded Data Parallelism）功能的具体指南。

文档链接:
- https://github.com/NVIDIA/Megatron-LM/blob/main/docs/source/api-guide/custom_fsdp.md
如何在 Megatron-Core 中启用自定义 FSDP？
- --use-custom-fsdp
- --data-parallel-sharding-strategy optim_grads_params
- --no-gradient-accumulation-fusion # 注：自定义FSDP（c-fsdp）目前不支持 TE（TransformerEngine）的梯度累积融合。
可选标志 (Optional Flags):
- --calculate-per-token-loss # 用于与普通数据并行获得更好的数值对齐。
- --init-model-with-meta-device # 在初始化巨大模型时避免 OOM（内存不足）的必要配置。